package mylab.starters.dag;

import com.google.common.collect.Sets;

import java.time.Duration;
import java.time.LocalDateTime;
import java.util.Arrays;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import mylab.base.foundation.lang.Callback;
import mylab.starters.dag.enums.DagState;
import mylab.starters.dag.graph.Edge;
import mylab.starters.dag.graph.Vertex;
import mylab.starters.dag.graph.impl.BaseGraph;
import mylab.starters.dag.task.Engine;
import mylab.starters.dag.task.EngineCallback;
import mylab.starters.dag.task.Task;
import mylab.utils.common.StringUtil;
import mylab.utils.common.concurrent.Concurrent;

@Slf4j
@SuppressWarnings("all")
public class DagEngine extends BaseGraph implements Engine {

    DagContext context;
    DagState state;
    Builder builder;

    // 全局中断标识：true 表示需中止所有任务
    AtomicBoolean atomicAbort;

    @Getter
    EngineCallback engineCallback;

    private DagEngine(Builder builder) {
        this.builder = builder;
        if (StringUtil.isNotEmpty(builder.id)) {
            this.id = builder.id;
        }
        if (StringUtil.isNotEmpty(builder.name)) {
            this.name = builder.name;
        }

        addVertices(builder.vertices);
        addEdges(builder.edges);
    }

    @Override
    public boolean isAborted() {
        return atomicAbort != null ? atomicAbort.get() : false;
    }

    @Override
    public void abort(boolean interrupt) {
        if (!Objects.equals(DagState.RUNNING, state)) {
            return;
        }

        if (atomicAbort != null) {
            atomicAbort.compareAndSet(false, true);
        }

        if (interrupt && this.context != null) {
            this.context.interruptThread();
        }
    }

    @Override
    public void syncRun(long timeout, TimeUnit timeUnit, EngineCallback callback) {
        init();
        this.engineCallback = callback;
        Runnable runnable = getDagInitialTask(getBeginVertices(), timeout, timeUnit, callback);
        runnable.run();
    }

    @Override
    public void asyncRun(long timeout, TimeUnit timeUnit, EngineCallback callback) {
        init();
        this.engineCallback = callback;
        Runnable runnable = getDagInitialTask(getBeginVertices(), timeout, timeUnit, callback);
        context.getExecutorService().submit(runnable);
    }

    private void init() {
        this.state = DagState.INIT;
        this.atomicAbort = new AtomicBoolean(false);
        this.context = new DagContext(this);
        this.context.setExecutorService(builder.executorService == null ? Concurrent.executor : builder.executorService);

        vertices.forEach(vertex -> {
            vertex.getIndegree().set(vertex.getDependentVertices().size());
            if (vertex instanceof Task<?> task) {
                task.isExecuted().set(false);
            }
        });

        if (engineCallback != null) {
            engineCallback.onEngineBegin(getId());
        }
    }

    private Runnable getDagInitialTask(Set<Vertex> beginTasks, long timeout, TimeUnit timeUnit, EngineCallback callback) {

        check();

        return () -> {
            if (Objects.equals(DagState.RUNNING, state)) {
                log.warn("DAG engine is running, please try again later");
                return;
            }
            try {
                state = DagState.RUNNING;
                LocalDateTime begin = LocalDateTime.now();
                DagContextHolder.set(context);
                beginTasks.forEach(vertex -> {
                    if (vertex instanceof Task<?> task) {
                        task.doRun();
                    }
                });

                boolean isTimeout = false;
                try {
                    isTimeout = !context.getSyncLatch().await(timeout, timeUnit);
                } catch (InterruptedException e) {
                    log.error("DAG Engine is interrupted", e);
                }

                if (isTimeout) {
                    state = DagState.TIMEOUT;
                    abort(true);
                    Callback.dealOnFailure(callback, StringUtil.format("任务超时。id:{} timeout:{} unit:{}", getId(), timeout, timeUnit));
                } else {
                    state = DagState.FINISH;
                    LocalDateTime end = LocalDateTime.now();
                    log.info("DAG 执行耗时:{}秒", Duration.between(begin, end).toSeconds());

                    DagResult dagResult = new DagResult(getId(), context.getResultMap());
                    Callback.dealOnSuccess(callback, dagResult);
                }
            } catch (Exception e) {
                state = DagState.ERROR;
                Callback.dealOnException(callback, e);
            } finally {
                DagContextHolder.remove(getId());
                if (engineCallback != null) {
                    engineCallback.onEnginEnd(getId(), context.getResultMap());
                }
            }
        };
    }


    public static class Builder {
        String id;
        String name;
        ExecutorService executorService;
        Set<Vertex> vertices = Sets.newHashSet();
        Set<Edge> edges = Sets.newHashSet();

        public Builder id(String id) {
            this.id = id;
            return this;
        }

        public Builder name(String name) {
            this.name = name;
            return this;
        }

        public Builder executorService(ExecutorService executorService) {
            this.executorService = executorService;
            return this;
        }

        public Builder addVertices(Vertex... vertices) {
            // Android 设备的运行时环境（ART）可能只支持到 Java 8 的 Stream API，而 Stream.toList() 是 Java 16 的特性。
            //this.vertices.addAll(Arrays.stream(vertices).toList());
            this.vertices.addAll(Arrays.stream(vertices).collect(Collectors.toList()));
            return this;
        }

        public Builder addEdges(Edge... edges) {
            this.edges.addAll(Arrays.stream(edges).toList());
            return this;
        }

        public DagEngine build() {
            return new DagEngine(this);
        }

    }
}
